# views.py
import base64
import time

import cv2
import numpy as np
import uuid
import os
from ultralytics import YOLO
from django.http import JsonResponse, StreamingHttpResponse, HttpResponseNotFound
from django.conf import settings

# YOLO 模型加载
model = YOLO("D:\\Python\\project\\ai-training\\expression_best.pt")
# todo 改路径，改成.pt在自己工程项目中的路径
# 李煊豪的路径model = YOLO("V:\pycharm2024.1.4\py_workspace\AI_FaceDetection\expression_best.pt")
# 用一个字典记录用户上传的视频路径，方便流式播放
VIDEO_CACHE = {}

def predict(request):
    """通用预测接口：接收图片或视频，并返回处理结果"""
    if request.method != "POST":
        return JsonResponse({"code": 405, "msg": "只支持 POST"})

    if "file" not in request.FILES:
        return JsonResponse({"code": 400, "msg": "没有找到文件"})

    uploaded_file = request.FILES["file"]
    content_type = uploaded_file.content_type

    # === 图片处理 ===
    if content_type.startswith("image"):
        img_bytes = uploaded_file.read()
        img_array = np.frombuffer(img_bytes, np.uint8)
        img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)

        results = model(img)
        result_img_array = results[0].plot()

        _, encoded_img = cv2.imencode('.jpg', result_img_array)
        img_base64 = base64.b64encode(encoded_img.tobytes()).decode('utf-8')

        return JsonResponse({
            "code": 200,
            "type": "image",
            "msg": "图片识别成功",
            "processed_img_base64": img_base64
        })

    # === 视频处理 ===
    elif content_type.startswith("video"):
        # 用 UUID 生成唯一文件名
        file_id = uuid.uuid4().hex
        temp_filename = f"{file_id}.mp4"
        temp_path = os.path.join(settings.BASE_DIR, "static", "videos", temp_filename)
        os.makedirs(os.path.dirname(temp_path), exist_ok=True) #如果路径不存在，创建
        # 保存视频文件
        with open(temp_path, "wb") as f:
            for chunk in uploaded_file.chunks():
                f.write(chunk)

        # 缓存路径，供 video_stream 调用
        VIDEO_CACHE[file_id] = temp_path

        return JsonResponse({
            "code": 200,
            "type": "video",
            "msg": "视频已接收",
            "stream_url": f"/video_stream/{file_id}"  # 前端直接访问此URL播放流
        })

    return JsonResponse({"code": 415, "msg": "不支持的文件类型"})


def video_stream(request, file_id):
    """专门用于视频流的接口"""
    if file_id not in VIDEO_CACHE:
        return HttpResponseNotFound("视频不存在或已被删除")

    video_path = VIDEO_CACHE[file_id]

    def generate():
        cap = cv2.VideoCapture(video_path)

        # 读取原视频 FPS，避免播放过快
        fps = cap.get(cv2.CAP_PROP_FPS)
        if fps <= 0:
            fps = 25  # 如果无法读取 fps，设个默认值
        frame_interval = 1.0 / fps

        while cap.isOpened():
            start_time = time.time()

            ret, frame = cap.read()
            if not ret:
                break
            results = model(frame)
            result_frame = results[0].plot()

            _, buffer = cv2.imencode('.jpg', result_frame)
            frame_bytes = buffer.tobytes()

            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + frame_bytes + b'\r\n')

            # 保持帧率：计算本帧推理耗时后睡眠剩余时间
            elapsed = time.time() - start_time
            if elapsed < frame_interval:
                time.sleep(frame_interval - elapsed)

        cap.release()
        # 推理完成后清理临时文件，释放存储空间
        if os.path.exists(video_path):
            os.remove(video_path)
        VIDEO_CACHE.pop(file_id, None)

    return StreamingHttpResponse(
        generate(),
        content_type='multipart/x-mixed-replace; boundary=frame'
    )
